package cn.kinoko.common.utils;

import lombok.extern.slf4j.Slf4j;

import javax.crypto.Cipher;
import java.io.ByteArrayOutputStream;
import java.io.FileNotFoundException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;

@Slf4j
public class RSAUtil {

    /**
     * 密钥类实例化入参
     */
    private static final String KEY_ALGORITHM = "RSA";

    private static final String PUBLIC_KEY = "RSAPublicKey";
    private static final String PRIVATE_KEY = "RSAPrivateKey";

    /**
     * RSA最大加密明文大小
     */
    private static final int MAX_ENCRYPT_BLOCK = 245;

    /**
     * RSA最大解密密文大小，
     * RSA 位数 如果采用1024 上面最大加密和最大解密则须填写: 117 128
     * RSA 位数 如果采用2048 上面最大加密和最大解密则须填写: 245 256
     */
    private static final int MAX_DECRYPT_BLOCK = 256;

    /**
     * 初始化密钥对生成器时，指定密钥大小的整数值（安全漏洞，长度至少为2048）
     */
    private static final int KEY_PAIR_INIT_SIZE = 2048;

    /**
     * 从文件中读取公钥
     * @param path 公钥保存路径，相对于classpath
     * @return 公钥对象
     */
    public static PublicKey readPublicKey(String path) throws Exception {
        byte[] bytes = readFile(path);
        return getPublicKey(bytes);
    }

    /**
     * 从文件中读取密钥
     * @param path 私钥保存路径，相对于classpath
     * @return 私钥对象
     */
    public static PrivateKey readPrivateKey(String path) throws Exception {
        byte[] bytes = readFile(path);
        return getPrivateKey(bytes);
    }

    /**
     * 生成密钥对
     * @return 密钥对
     * @throws Exception e
     */
    private static Map<String, Object> initKey() throws Exception {
        KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance(KEY_ALGORITHM);
        keyPairGen.initialize(KEY_PAIR_INIT_SIZE);
        KeyPair keyPair = keyPairGen.generateKeyPair();
        RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();
        RSAPrivateKey privateKey = (RSAPrivateKey) keyPair.getPrivate();
        Map<String, Object> keyMap = new HashMap<>(2);
        keyMap.put(PUBLIC_KEY, publicKey);
        keyMap.put(PRIVATE_KEY, privateKey);
        return keyMap;
    }

    /**
     * 获取公钥字符串
     * @param keyMap 密钥对
     * @return 公钥字符串
     */
    private static String getPublicKeyStr(Map<String, Object> keyMap) {
        // 获得 map 中的公钥对象，转为 key 对象
        Key key = (Key) keyMap.get(PUBLIC_KEY);
        // 编码返回字符串
        return encryptBASE64(key.getEncoded());
    }

    /**
     * 获取私钥字符串
     * @param keyMap 密钥对
     * @return 获取私钥字符串
     */
    private static String getPrivateKeyStr(Map<String, Object> keyMap) {
        // 获得 map 中的私钥对象，转为 key 对象
        Key key = (Key) keyMap.get(PRIVATE_KEY);
        // 编码返回字符串
        return encryptBASE64(key.getEncoded());
    }

    /**
     * 获取私钥字符串
     * @param key 密钥
     * @return 获取密钥字符串
     */
    private static String getKeyStr(Key key) {
        // 编码返回字符串
        return encryptBASE64(key.getEncoded());
    }

    /**
     * 获取公钥
     * @param publicKeyString 公钥编码后的字符串
     * @return 公钥
     */
    private static PublicKey getPublicKey(String publicKeyString) throws NoSuchAlgorithmException, InvalidKeySpecException {
        byte[] publicKeyByte = Base64.getMimeDecoder().decode(publicKeyString);
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(publicKeyByte);
        KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
        return keyFactory.generatePublic(keySpec);
    }

    /**
     * 获取私钥
     * @param privateKeyString 私钥编码后的字符串
     * @return 私钥
     * @throws Exception e
     */
    private static PrivateKey getPrivateKey(String privateKeyString) throws Exception {
        byte[] privateKeyByte = Base64.getMimeDecoder().decode(privateKeyString);
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(privateKeyByte);
        KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
        return keyFactory.generatePrivate(keySpec);
    }

    /**
     * 从文件读取
     * @param filePath 文件路径
     * @return 文件字节
     */
    private static byte[] readFile(String filePath) throws Exception {
        try (InputStream is = RSAUtil.class.getClassLoader().getResourceAsStream(filePath)) {
            if (is == null) {
                throw new FileNotFoundException(filePath + " not found");
            }
            return is.readAllBytes();
        }
    }

    /**
     * 获取公钥
     * @param bytes 公钥的字节形式
     * @return 公钥
     */
    private static PublicKey getPublicKey(byte[] bytes) throws Exception {
        bytes = Base64.getMimeDecoder().decode(bytes);
        X509EncodedKeySpec spec = new X509EncodedKeySpec(bytes);
        KeyFactory factory = KeyFactory.getInstance(KEY_ALGORITHM);
        return factory.generatePublic(spec);
    }

    /**
     * 获取私钥
     * @param bytes 私钥的字节形式
     * @return 私钥
     */
    private static PrivateKey getPrivateKey(byte[] bytes) throws NoSuchAlgorithmException, InvalidKeySpecException {
        bytes = Base64.getMimeDecoder().decode(bytes);
        PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(bytes);
        KeyFactory factory = KeyFactory.getInstance(KEY_ALGORITHM);
        return factory.generatePrivate(spec);
    }

    /**
     * BASE64 编码返回加密字符串
     * @param key 需要编码的字节数组
     * @return 编码后的字符串
     */
    private static String encryptBASE64(byte[] key) {
        return Base64.getMimeEncoder().encodeToString(key);
    }

    /**
     * 公钥加密
     * @param text         待加密的明文字符串
     * @param publicKeyStr 公钥
     * @return 加密后的密文
     */
    private static String encrypt(String text, String publicKeyStr) {
        try {
            log.info("明文字符串为:[{}]", text);
            Cipher cipher = Cipher.getInstance(KEY_ALGORITHM);
            cipher.init(Cipher.ENCRYPT_MODE, getPublicKey(publicKeyStr));
            byte[] tempBytes = cipher.doFinal(text.getBytes(StandardCharsets.UTF_8));
            return Base64.getEncoder().encodeToString(tempBytes);
        } catch (Exception e) {
            throw new RuntimeException("加密字符串[" + text + "]时遇到异常", e);
        }
    }

    /**
     * 私钥解密
     * @param secretText    待解密的密文字符串
     * @param privateKeyStr 私钥
     * @return 解密后的明文
     */
    private static String decrypt(String secretText, String privateKeyStr) {
        try {
            // 生成私钥
            Cipher cipher = Cipher.getInstance(KEY_ALGORITHM);
            cipher.init(Cipher.DECRYPT_MODE, getPrivateKey(privateKeyStr));
            // 密文解码
            byte[] secretTextDecoded = Base64.getDecoder().decode(secretText.getBytes(StandardCharsets.UTF_8));
            byte[] tempBytes = cipher.doFinal(secretTextDecoded);
            return new String(tempBytes);
        } catch (Exception e) {
            throw new RuntimeException("解密字符串[" + secretText + "]时遇到异常", e);
        }
    }

    /**
     * 公钥分段加密
     * @param plainText    待加密的明文字符串
     * @param publicKeyStr 公钥
     * @return 加密后的密文
     */
    private static String sectionEncrypt(String plainText, String publicKeyStr) throws Exception {
        byte[] plainTextArray = plainText.getBytes(StandardCharsets.UTF_8);
        PublicKey publicKey = getPublicKey(publicKeyStr);
        Cipher cipher = Cipher.getInstance(KEY_ALGORITHM);
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        int inputLen = plainTextArray.length;
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        int offSet = 0;
        int i = 0;
        byte[] cache;
        while (inputLen - offSet > 0) {
            if (inputLen - offSet > MAX_ENCRYPT_BLOCK) {
                cache = cipher.doFinal(plainTextArray, offSet, MAX_ENCRYPT_BLOCK);
            } else {
                cache = cipher.doFinal(plainTextArray, offSet, inputLen - offSet);
            }
            out.write(cache, 0, cache.length);
            i++;
            offSet = i * MAX_ENCRYPT_BLOCK;
        }
        byte[] encryptText = out.toByteArray();
        out.close();
        return Base64.getEncoder().encodeToString(encryptText);
    }

    /**
     * 私钥分段解密
     * @param encryptTextHex 待解密的密文字符串
     * @param privateKeyStr  私钥
     * @return 解密后的明文
     */
    private static String sectionDecrypt(String encryptTextHex, String privateKeyStr) throws Exception {
        byte[] encryptText = Base64.getDecoder().decode(encryptTextHex);
        PrivateKey privateKey = getPrivateKey(privateKeyStr);
        Cipher cipher = Cipher.getInstance(KEY_ALGORITHM);
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        int inputLen = encryptText.length;
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        int offSet = 0;
        byte[] cache;
        int i = 0;
        // 对数据分段解密
        while (inputLen - offSet > 0) {
            if (inputLen - offSet > MAX_DECRYPT_BLOCK) {
                cache = cipher.doFinal(encryptText, offSet, MAX_DECRYPT_BLOCK);
            } else {
                cache = cipher.doFinal(encryptText, offSet, inputLen - offSet);
            }
            out.write(cache, 0, cache.length);
            i++;
            offSet = i * MAX_DECRYPT_BLOCK;
        }
        out.close();
        return out.toString();
    }

    public static void main(String[] args) throws Exception {
        String cipherText;
        // 原始明文
        String content = "春江潮水连海平，海上明月共潮生。滟滟随波千万里，何处春江无月明。";

        // 生成密钥对
        Map<String, Object> keyMap = initKey();
        String publicKey = getKeyStr((Key) keyMap.get(PUBLIC_KEY));
        // 打印公钥、私钥
        System.out.println("公钥：（填充方式：PKCS1_PADDING，输出类型：base64，字符集：utf8编码）");
        System.out.println("-----BEGIN PUBLIC KEY-----");
        System.out.println(publicKey);
        System.out.println("-----END PUBLIC KEY-----");
        System.out.println("\n");
        String privateKey = getKeyStr((Key) keyMap.get(PRIVATE_KEY));
        System.out.println("私钥：（填充方式：PKCS1_PADDING，输出类型：base64，字符集：utf8编码）");
        System.out.println("-----BEGIN RSA PRIVATE KEY-----");
        System.out.println(privateKey);
        System.out.println("-----END RSA PRIVATE KEY-----");
        System.out.println("\n");

        // 加密
        cipherText = sectionEncrypt(content, publicKey);
        log.info("加密后的密文:[{}]，长度:[{}]", cipherText, cipherText.length());

        // 解密
        String plainText = sectionDecrypt(cipherText, privateKey);
        log.info("解密后明文:[{}]", plainText);

    }
}